package com.hyz.yuaiagent.rag.query;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 多查询扩展器演示类
 * 该类演示了如何使用MultiQueryExpander来扩展查询
 */
@Component
public class MultiQueryExpanderDemo {

    // DashScope聊天模型实例，用于创建ChatClient
    private final DashScopeChatModel dashscopeChatModel;

    /**
     * 构造函数
     *
     * @param dashscopeChatModel DashScope聊天模型，用于初始化ChatClient
     */
    public MultiQueryExpanderDemo(DashScopeChatModel dashscopeChatModel) {
        this.dashscopeChatModel = dashscopeChatModel;
    }

    /**
     * 扩展查询
     * 使用MultiQueryExpander将原始查询扩展为多个相关的查询
     *
     * @param query 原始查询
     * @return 扩展后的查询列表
     */
    public List<Query> expand(Query query) {
        // 创建MultiQueryExpander实例，配置ChatClient和生成查询的数量
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(ChatClient.builder(dashscopeChatModel))
                .numberOfQueries(3)
                .build();

        // 使用MultiQueryExpander扩展原始查询并返回扩展后的查询列表
        return queryExpander.expand(query);
    }

}
